{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Closed-Form Policy to Play Pendulum-v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import logging\n",
    "import itertools\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import gym\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "        format='%(asctime)s [%(levelname)s] %(message)s',\n",
    "        stream=sys.stdout, datefmt='%H:%M:%S')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:00:00 [INFO] id: Pendulum-v1\n",
      "00:00:00 [INFO] entry_point: gym.envs.classic_control:PendulumEnv\n",
      "00:00:00 [INFO] reward_threshold: None\n",
      "00:00:00 [INFO] nondeterministic: False\n",
      "00:00:00 [INFO] max_episode_steps: 200\n",
      "00:00:00 [INFO] order_enforce: True\n",
      "00:00:00 [INFO] _kwargs: {}\n",
      "00:00:00 [INFO] _env_name: Pendulum\n",
      "00:00:00 [INFO] max_speed: 8\n",
      "00:00:00 [INFO] max_torque: 2.0\n",
      "00:00:00 [INFO] dt: 0.05\n",
      "00:00:00 [INFO] g: 10.0\n",
      "00:00:00 [INFO] m: 1.0\n",
      "00:00:00 [INFO] l: 1.0\n",
      "00:00:00 [INFO] viewer: None\n",
      "00:00:00 [INFO] action_space: Box([-2.], [2.], (1,), float32)\n",
      "00:00:00 [INFO] observation_space: Box([-1. -1. -8.], [1. 1. 8.], (3,), float32)\n",
      "00:00:00 [INFO] np_random: RandomState(MT19937)\n",
      "00:00:00 [INFO] spec: EnvSpec(Pendulum-v1)\n"
     ]
    }
   ],
   "source": [
    "env = gym.make('Pendulum-v1')\n",
    "env.seed(0)\n",
    "for key in vars(env.spec):\n",
    "    logging.info('%s: %s', key, vars(env.spec)[key])\n",
    "for key in vars(env.unwrapped):\n",
    "    logging.info('%s: %s', key, vars(env.unwrapped)[key])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClosedFormAgent:\n",
    "    def __init__(self, _):\n",
    "        pass\n",
    "\n",
    "    def reset(self, mode=None):\n",
    "        pass\n",
    "\n",
    "    def step(self, observation, _reward, _done):\n",
    "        x, y, angle_velocity = observation\n",
    "        flip = (y < 0.)\n",
    "        if flip:\n",
    "            y *= -1. # now y >= 0\n",
    "            angle_velocity *= -1.\n",
    "        angle = np.arcsin(y)\n",
    "        if x < 0.:\n",
    "            angle = np.pi - angle\n",
    "        if (angle < -0.3 * angle_velocity) or \\\n",
    "                (angle > 0.03 * (angle_velocity - 2.5) ** 2. + 1. and \\\n",
    "                angle < 0.15 * (angle_velocity + 3.) ** 2. + 2.):\n",
    "            force = 2.\n",
    "        else:\n",
    "            force = -2.\n",
    "        if flip:\n",
    "            force *= -1.\n",
    "        action = np.array([force,])\n",
    "        return action\n",
    "\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "\n",
    "agent = ClosedFormAgent(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "00:00:00 [INFO] ==== test ====\n",
      "00:00:00 [DEBUG] test episode 0: reward = -267.55, steps = 200\n",
      "00:00:00 [DEBUG] test episode 1: reward = -126.17, steps = 200\n",
      "00:00:00 [DEBUG] test episode 2: reward = -235.76, steps = 200\n",
      "00:00:00 [DEBUG] test episode 3: reward = -127.66, steps = 200\n",
      "00:00:00 [DEBUG] test episode 4: reward = -231.72, steps = 200\n",
      "00:00:00 [DEBUG] test episode 5: reward = -283.73, steps = 200\n",
      "00:00:01 [DEBUG] test episode 6: reward = -2.03, steps = 200\n",
      "00:00:01 [DEBUG] test episode 7: reward = -225.50, steps = 200\n",
      "00:00:01 [DEBUG] test episode 8: reward = -2.29, steps = 200\n",
      "00:00:01 [DEBUG] test episode 9: reward = -283.62, steps = 200\n",
      "00:00:01 [DEBUG] test episode 10: reward = -120.88, steps = 200\n",
      "00:00:01 [DEBUG] test episode 11: reward = -3.70, steps = 200\n",
      "00:00:01 [DEBUG] test episode 12: reward = -238.70, steps = 200\n",
      "00:00:01 [DEBUG] test episode 13: reward = -119.87, steps = 200\n",
      "00:00:01 [DEBUG] test episode 14: reward = -1.69, steps = 200\n",
      "00:00:01 [DEBUG] test episode 15: reward = -120.17, steps = 200\n",
      "00:00:01 [DEBUG] test episode 16: reward = -2.05, steps = 200\n",
      "00:00:01 [DEBUG] test episode 17: reward = -126.85, steps = 200\n",
      "00:00:01 [DEBUG] test episode 18: reward = -230.96, steps = 200\n",
      "00:00:01 [DEBUG] test episode 19: reward = -1.69, steps = 200\n",
      "00:00:01 [DEBUG] test episode 20: reward = -127.73, steps = 200\n",
      "00:00:01 [DEBUG] test episode 21: reward = -338.72, steps = 200\n",
      "00:00:01 [DEBUG] test episode 22: reward = -301.22, steps = 200\n",
      "00:00:01 [DEBUG] test episode 23: reward = -221.75, steps = 200\n",
      "00:00:01 [DEBUG] test episode 24: reward = -127.78, steps = 200\n",
      "00:00:01 [DEBUG] test episode 25: reward = -221.79, steps = 200\n",
      "00:00:01 [DEBUG] test episode 26: reward = -237.35, steps = 200\n",
      "00:00:01 [DEBUG] test episode 27: reward = -120.55, steps = 200\n",
      "00:00:01 [DEBUG] test episode 28: reward = -119.65, steps = 200\n",
      "00:00:01 [DEBUG] test episode 29: reward = -336.56, steps = 200\n",
      "00:00:01 [DEBUG] test episode 30: reward = -245.63, steps = 200\n",
      "00:00:01 [DEBUG] test episode 31: reward = -238.51, steps = 200\n",
      "00:00:01 [DEBUG] test episode 32: reward = -3.20, steps = 200\n",
      "00:00:01 [DEBUG] test episode 33: reward = -118.62, steps = 200\n",
      "00:00:01 [DEBUG] test episode 34: reward = -223.51, steps = 200\n",
      "00:00:01 [DEBUG] test episode 35: reward = -232.35, steps = 200\n",
      "00:00:01 [DEBUG] test episode 36: reward = -352.74, steps = 200\n",
      "00:00:01 [DEBUG] test episode 37: reward = -4.81, steps = 200\n",
      "00:00:01 [DEBUG] test episode 38: reward = -124.59, steps = 200\n",
      "00:00:01 [DEBUG] test episode 39: reward = -116.71, steps = 200\n",
      "00:00:02 [DEBUG] test episode 40: reward = -240.21, steps = 200\n",
      "00:00:02 [DEBUG] test episode 41: reward = -349.29, steps = 200\n",
      "00:00:02 [DEBUG] test episode 42: reward = -124.80, steps = 200\n",
      "00:00:02 [DEBUG] test episode 43: reward = -2.00, steps = 200\n",
      "00:00:02 [DEBUG] test episode 44: reward = -127.37, steps = 200\n",
      "00:00:02 [DEBUG] test episode 45: reward = -120.40, steps = 200\n",
      "00:00:02 [DEBUG] test episode 46: reward = -116.89, steps = 200\n",
      "00:00:02 [DEBUG] test episode 47: reward = -120.74, steps = 200\n",
      "00:00:02 [DEBUG] test episode 48: reward = -3.17, steps = 200\n",
      "00:00:02 [DEBUG] test episode 49: reward = -1.92, steps = 200\n",
      "00:00:02 [DEBUG] test episode 50: reward = -225.52, steps = 200\n",
      "00:00:02 [DEBUG] test episode 51: reward = -123.44, steps = 200\n",
      "00:00:02 [DEBUG] test episode 52: reward = -116.63, steps = 200\n",
      "00:00:02 [DEBUG] test episode 53: reward = -119.64, steps = 200\n",
      "00:00:02 [DEBUG] test episode 54: reward = -120.17, steps = 200\n",
      "00:00:02 [DEBUG] test episode 55: reward = -229.39, steps = 200\n",
      "00:00:02 [DEBUG] test episode 56: reward = -127.65, steps = 200\n",
      "00:00:02 [DEBUG] test episode 57: reward = -122.00, steps = 200\n",
      "00:00:02 [DEBUG] test episode 58: reward = -118.49, steps = 200\n",
      "00:00:02 [DEBUG] test episode 59: reward = -115.77, steps = 200\n",
      "00:00:02 [DEBUG] test episode 60: reward = -126.50, steps = 200\n",
      "00:00:02 [DEBUG] test episode 61: reward = -119.58, steps = 200\n",
      "00:00:02 [DEBUG] test episode 62: reward = -115.18, steps = 200\n",
      "00:00:02 [DEBUG] test episode 63: reward = -118.86, steps = 200\n",
      "00:00:02 [DEBUG] test episode 64: reward = -115.36, steps = 200\n",
      "00:00:02 [DEBUG] test episode 65: reward = -123.53, steps = 200\n",
      "00:00:02 [DEBUG] test episode 66: reward = -2.37, steps = 200\n",
      "00:00:02 [DEBUG] test episode 67: reward = -239.47, steps = 200\n",
      "00:00:02 [DEBUG] test episode 68: reward = -222.48, steps = 200\n",
      "00:00:02 [DEBUG] test episode 69: reward = -117.72, steps = 200\n",
      "00:00:02 [DEBUG] test episode 70: reward = -121.58, steps = 200\n",
      "00:00:02 [DEBUG] test episode 71: reward = -128.51, steps = 200\n",
      "00:00:03 [DEBUG] test episode 72: reward = -218.31, steps = 200\n",
      "00:00:03 [DEBUG] test episode 73: reward = -2.90, steps = 200\n",
      "00:00:03 [DEBUG] test episode 74: reward = -115.21, steps = 200\n",
      "00:00:03 [DEBUG] test episode 75: reward = -3.35, steps = 200\n",
      "00:00:03 [DEBUG] test episode 76: reward = -116.52, steps = 200\n",
      "00:00:03 [DEBUG] test episode 77: reward = -125.39, steps = 200\n",
      "00:00:03 [DEBUG] test episode 78: reward = -118.11, steps = 200\n",
      "00:00:03 [DEBUG] test episode 79: reward = -127.57, steps = 200\n",
      "00:00:03 [DEBUG] test episode 80: reward = -244.67, steps = 200\n",
      "00:00:03 [DEBUG] test episode 81: reward = -3.62, steps = 200\n",
      "00:00:03 [DEBUG] test episode 82: reward = -337.73, steps = 200\n",
      "00:00:03 [DEBUG] test episode 83: reward = -133.64, steps = 200\n",
      "00:00:03 [DEBUG] test episode 84: reward = -122.92, steps = 200\n",
      "00:00:03 [DEBUG] test episode 85: reward = -124.27, steps = 200\n",
      "00:00:03 [DEBUG] test episode 86: reward = -339.02, steps = 200\n",
      "00:00:03 [DEBUG] test episode 87: reward = -117.75, steps = 200\n",
      "00:00:03 [DEBUG] test episode 88: reward = -127.78, steps = 200\n",
      "00:00:03 [DEBUG] test episode 89: reward = -118.56, steps = 200\n",
      "00:00:03 [DEBUG] test episode 90: reward = -124.17, steps = 200\n",
      "00:00:03 [DEBUG] test episode 91: reward = -120.17, steps = 200\n",
      "00:00:03 [DEBUG] test episode 92: reward = -116.71, steps = 200\n",
      "00:00:03 [DEBUG] test episode 93: reward = -127.08, steps = 200\n",
      "00:00:03 [DEBUG] test episode 94: reward = -2.01, steps = 200\n",
      "00:00:03 [DEBUG] test episode 95: reward = -238.65, steps = 200\n",
      "00:00:03 [DEBUG] test episode 96: reward = -124.64, steps = 200\n",
      "00:00:03 [DEBUG] test episode 97: reward = -118.74, steps = 200\n",
      "00:00:03 [DEBUG] test episode 98: reward = -120.61, steps = 200\n",
      "00:00:03 [DEBUG] test episode 99: reward = -219.02, steps = 200\n",
      "00:00:03 [INFO] average episode reward = -145.52 ± 90.77\n"
     ]
    }
   ],
   "source": [
    "def play_episode(env, agent, max_episode_steps=None, mode=None, render=False):\n",
    "    observation, reward, done = env.reset(), 0., False\n",
    "    agent.reset(mode=mode)\n",
    "    episode_reward, elapsed_steps = 0., 0\n",
    "    while True:\n",
    "        action = agent.step(observation, reward, done)\n",
    "        if render:\n",
    "            env.render()\n",
    "        if done:\n",
    "            break\n",
    "        observation, reward, done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        elapsed_steps += 1\n",
    "        if max_episode_steps and elapsed_steps >= max_episode_steps:\n",
    "            break\n",
    "    agent.close()\n",
    "    return episode_reward, elapsed_steps\n",
    "\n",
    "\n",
    "logging.info('==== test ====')\n",
    "episode_rewards = []\n",
    "for episode in range(100):\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('test episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "logging.info('average episode reward = %.2f ± %.2f',\n",
    "        np.mean(episode_rewards), np.std(episode_rewards))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
